using System;
using Diversity.Maths;

namespace Diversity.Robotics.Navigation
{
    /// <summary>
    /// A utility class for helping with a simple extended kalman filter.
    /// </summary>
    public class ExtendedKalmanFilter
    {
        private double _alpha1;
        private double _alpha2;
        private double _alpha3;
        private double _alpha4;
        private double _observationDistanceVariance;
        private double _observationAngleVariance;

        /// <summary>
        /// Motion model error 1
        /// </summary>
        public double Alpha1
        {
            get { return _alpha1; }
            set { _alpha1 = value; }
        }

        /// <summary>
        /// Motion model error 2
        /// </summary>
        public double Alpha2
        {
            get { return _alpha2; }
            set { _alpha2 = value; }
        }

        /// <summary>
        /// Motion model error 3
        /// </summary>
        public double Alpha3
        {
            get { return _alpha3; }
            set { _alpha3 = value; }
        }

        /// <summary>
        /// Motion model error 4
        /// </summary>
        public double Alpha4
        {
            get { return _alpha4; }
            set { _alpha4 = value; }
        }

       
        /// <summary>
        /// Sets the four error parameters of the motion model.
        /// </summary>
        /// <param name="alpha1">Alpha 1</param>
        /// <param name="alpha2">Alpha 2</param>
        /// <param name="alpha3">Alpha 3</param>
        /// <param name="alpha4">Alpha 4</param>
        public void SetMotionErrors(double alpha1, double alpha2, double alpha3, double alpha4)
        {
            _alpha1 = alpha1;
            _alpha2 = alpha2;
            _alpha3 = alpha3;
            _alpha4 = alpha4;
        }

        /// <summary>
        /// Sets the two observation errors.
        /// </summary>
        /// <param name="distanceVariance">The variance in the observed distance.</param>
        /// <param name="angleVariance">The variance in the observed angle.</param>
        public void SetObservationErrors(double distanceVariance, double angleVariance)
        {
            _observationDistanceVariance = distanceVariance;
            _observationAngleVariance = angleVariance;
        }

        /// <summary>
        /// Just handles states, velocities and landmarks. 
        /// TODO: Should be generalised to bigger matrices
        /// </summary>
        public ExtendedKalmanFilter()
        {
        }

        /// <summary>
        /// The motion update where the new velocity is applied and uncertainty increases.
        /// </summary>
        /// <param name="state">The current state.</param>
        /// <param name="uncertainty">The current uncertainty matrix.</param>
        /// <param name="vel">The desired velocity.</param>
        /// <param name="t">The time for which the velocity is applied.</param>
        public void MotionUpdate(ref Pose2D state, ref Matrix uncertainty, Velocity vel, TimeSpan t)
        {
            // time of action
            double dTime = t.TotalSeconds;

            double velOverRot = (vel.Translation / vel.Rotation);
            double rotTime = vel.Rotation * dTime;
            double endTheta = state.Theta + rotTime;
            double cosTheta = Math.Cos(state.Theta);
            double sinTheta = Math.Sin(state.Theta);
            double cosEndTheta = Math.Cos(endTheta);
            double sinEndTheta = Math.Sin(endTheta);
            double tranSquared = vel.Translation * vel.Translation;
            double rotSquared = vel.Rotation * vel.Rotation;

            Matrix G = Matrix.Identity(3); // Projector of current state noise
            Matrix V = new Matrix(3, 2); // Projector of added control noise
            Matrix stateChange = new Matrix(3, 1); // True motion

            // special case for no rotation
            if (vel.Rotation == 0)
            {
                G[1, 2] = vel.Translation;

                V[0, 0] = -1;
                V[0, 1] = -1;
                V[1, 1] = vel.Translation / 2;
                V[2, 1] = dTime;

                stateChange[0, 0] = vel.Translation * dTime * cosTheta;
                stateChange[1, 0] = vel.Translation * dTime * sinTheta;
            }
            else
            {
                G[0, 2] = velOverRot * (-cosTheta + cosEndTheta);
                G[1, 2] = velOverRot * (-sinTheta + sinEndTheta);

                V[0, 0] = -(sinTheta + sinEndTheta) / vel.Rotation;
                V[0, 1] = (vel.Translation * (sinTheta - sinEndTheta) / rotSquared) +
                           (velOverRot * cosEndTheta * dTime);
                V[1, 0] = (cosTheta - cosEndTheta) / vel.Rotation;
                V[1, 1] = -(vel.Translation * (cosTheta - cosEndTheta) / rotSquared) +
                           (velOverRot * sinEndTheta * dTime);
                V[2, 1] = dTime;

                stateChange[0, 0] = velOverRot * (-sinTheta + sinEndTheta);
                stateChange[1, 0] = velOverRot * (cosTheta - cosEndTheta);
                stateChange[2, 0] = rotTime;
            }
            
            // Motion noise
            Matrix M = new Matrix(2, 2);
            M[0, 0] = _alpha1 * tranSquared + _alpha2 * rotSquared;
            M[1, 1] = _alpha3 * tranSquared + _alpha4 * rotSquared;

            // move the state
            state += stateChange;

            // calculate the uncertainty
            uncertainty = (G * uncertainty * ~G) + (V * M * ~V);
        }

        /// <summary>
        /// The sensor update where distance readings reduce the uncertainty.
        /// </summary>
        /// <param name="state">The current state.</param>
        /// <param name="uncertainty">The uncertainty matrix.</param>
        /// <param name="readings">The new readings.</param>
        /// <param name="map">An existing map of landmarks.</param>
        public void SensorUpdate(ref Pose2D state, ref Matrix uncertainty, SimpleDistanceReading[] readings, Landmark[] map)
        {
            // observation uncertainty
            Matrix Q = new Matrix(3, 3);
            Q[0, 0] = _observationDistanceVariance;
            Q[1, 1] = _observationAngleVariance;
            Q[2, 2] = 0.1; // meaningless

            for (int i = 0; i < readings.Length; i++)
            {
                // stuff it into a 3*1 matrix
                Matrix readingMatrix = new Matrix(3, 1);
                readingMatrix[0, 0] = readings[i].Distance;
                readingMatrix[1, 0] = readings[i].Theta;
                readingMatrix[2, 0] = i; // why not!

                double xDist = map[i].X - state.X;
                double yDist = map[i].Y - state.Y;
                double squaredDistance = Math.Pow(xDist, 2) + Math.Pow(yDist, 2);
                double dist = Math.Sqrt(squaredDistance);

                // do the same to the landmark
                Matrix mapMatrix = new Matrix(3, 1);
                mapMatrix[0, 0] = dist;
                mapMatrix[1, 0] = Math.Atan2(yDist,xDist);
                mapMatrix[2, 0] = i;

                // observation jacobian
                Matrix H = new Matrix(3, 3);
                H[0, 0] = -xDist / dist;
                H[0, 1] = -yDist / dist;
                H[1, 0] = yDist / squaredDistance;
                H[1, 1] = -xDist / squaredDistance;
                H[1, 2] = -1;
                Matrix HT = ~H;

                Matrix S = H * uncertainty * HT + Q;
                Matrix K = uncertainty * HT * !S;
                Matrix ObservationDisparity = readingMatrix - mapMatrix;

                // NEEDS CLEANING
                double angleDif = 0;
                if (readingMatrix[1,0] >= 0 && mapMatrix[1,0] >= 0)
                {
                    angleDif= readingMatrix[1, 0] - mapMatrix[1, 0];
                }
                else if (readingMatrix[1, 0] < 0 && mapMatrix[1, 0] < 0)
                {
                    angleDif= readingMatrix[1, 0] - mapMatrix[1, 0];
                }
                else
                {
                    angleDif= readingMatrix[1, 0] + mapMatrix[1, 0];
                    ObservationDisparity[1, 0] = angleDif;
                }
                if (angleDif != ObservationDisparity[1, 0])
                {
                    Console.Out.WriteLine("ahhh!");
                }

                Matrix stateChange = K*(ObservationDisparity);
                state += stateChange;
                uncertainty = (Matrix.Identity(3) - K*H)*uncertainty;
            }
        }

        /// <summary>
        /// Same as the other sensor update, but also finds likely landmarks.
        /// </summary>
        /// <param name="state">The current state.</param>
        /// <param name="uncertainty">The uncertainty matrix.</param>
        /// <param name="readings">The new readings.</param>
        /// <param name="map">An existing map of landmarks.</param>
        public void SensorUpdateUnknownCorrespondance(ref Pose2D state, ref Matrix uncertainty, SimpleDistanceReading[] readings, Landmark[] map)
        {
            // observation uncertainty
            Matrix Q = new Matrix(3, 3);
            Q[0, 0] = _observationDistanceVariance;
            Q[1, 1] = _observationAngleVariance;
            Q[2, 2] = 0.1; // meaningless

            for (int i = 0; i < readings.Length; i++)
            {
                int bestK = 0;
                Matrix BestInvertedS = new Matrix(3,3);
                Matrix BestH = new Matrix(3, 3);
                Matrix BestHT = new Matrix(3, 3);
                Matrix BestObservationDisparity = new Matrix(3,1);
                double maxProb = double.MaxValue;

                // stuff it into a 3*1 matrix
                Matrix readingMatrix = new Matrix(3, 1);
                readingMatrix[0, 0] = readings[i].Distance;
                readingMatrix[1, 0] = readings[i].Theta;
                readingMatrix[2, 0] = i; // why not!

                for (int k = 0; k < map.Length; k++)
                {
                 
                    // MUCH OF THIS SHOULD ONLY HAVE TO BE DONE ONCE
                    double xDist = map[k].X - state.X;
                    double yDist = map[k].Y - state.Y;
                    double squaredDistance = Math.Pow(xDist, 2) + Math.Pow(yDist, 2);
                    double dist = Math.Sqrt(squaredDistance);

                    // do the same to the landmark
                    Matrix mapMatrix = new Matrix(3, 1);
                    mapMatrix[0, 0] = dist;
                    mapMatrix[1, 0] = Math.Atan2(yDist, xDist);
                    mapMatrix[2, 0] = i;

                    // observation jacobian
                    Matrix H = new Matrix(3, 3);
                    H[0, 0] = -xDist / dist;
                    H[0, 1] = -yDist / dist;
                    H[1, 0] = yDist / squaredDistance;
                    H[1, 1] = -xDist / squaredDistance;
                    H[1, 2] = -1;
                    Matrix HT = ~H;
                    Matrix S = H * uncertainty * HT + Q;

                    Matrix ObservationDisparity = readingMatrix - mapMatrix;

                    // NEEDS CLEANING
                    double angleDif = 0;
                    if (readingMatrix[1, 0] >= 0 && mapMatrix[1, 0] >= 0)
                    {
                        angleDif = readingMatrix[1, 0] - mapMatrix[1, 0];
                    }
                    else if (readingMatrix[1, 0] < 0 && mapMatrix[1, 0] < 0)
                    {
                        angleDif = readingMatrix[1, 0] - mapMatrix[1, 0];
                    }
                    else
                    {
                        angleDif = readingMatrix[1, 0] + mapMatrix[1, 0];
                        ObservationDisparity[1, 0] = angleDif;
                    }
                    if (angleDif != ObservationDisparity[1, 0])
                    {
                        Console.Out.WriteLine("ahhh!");
                    }

                    Matrix InvertedS = !S;
                   // double detS = S.Determinant();
                    Matrix top = (~ObservationDisparity * InvertedS * ObservationDisparity); // reduces to 1*1 matrix
                    //Console.Out.WriteLine(top);
                    double prob = top[0, 0]; //(1.0 / (2 * Math.PI * detS)) * Math.Pow(Math.E, -1.0 / 2 * top[0,0]);

                    if (prob < maxProb)
                    {
                        bestK = k;
                        BestInvertedS = InvertedS;
                        BestObservationDisparity = ObservationDisparity;
                        BestH = H;
                        BestHT = HT;
                        maxProb = prob;
                    }
                  //  Console.Out.WriteLine("For " + i + " Prob of " + k + " is " + prob + " " + top[0,0]);
                }
              //  Console.Out.WriteLine("Best for " + i + " was " + bestK);
                Matrix K = uncertainty * BestHT * BestInvertedS;
               
                Matrix stateChange = K * (BestObservationDisparity);
                state += stateChange;
                uncertainty = (Matrix.Identity(3) - K * BestH) * uncertainty;
            }
        }
        
    }
}
